﻿#include "cipherutils.h"

#include <stdexcept>
#include <string>
#include <sstream>
#include <iomanip>

CipherUtils::CipherUtils() {
    // 初始化OpenSSL
    OpenSSL_add_all_algorithms();
    ERR_load_crypto_strings();
}
CipherUtils::~CipherUtils() {
    // 清理OpenSSL
    EVP_cleanup();
    ERR_free_strings();
}

CipherUtils& CipherUtils::getInstance() {
    static CipherUtils instance;
    return instance;
}

std::vector<std::shared_ptr<unsigned char[]>> CipherUtils::AES_generate() {
    // 生成256位的密码
    int keyLength = 256;
    std::shared_ptr<unsigned char[]> key(new unsigned char[keyLength]);
    std::shared_ptr<unsigned char[]> iv(new unsigned char[AES_BLOCK_SIZE]);

    if (!RAND_bytes(key.get(), keyLength)) {
        throw std::runtime_error("AES Key generate error");
    }

    if (!RAND_bytes(iv.get(), AES_BLOCK_SIZE)) {
        throw std::runtime_error("AES IV generate error");
    }

    return { key, iv };
}



std::pair<std::shared_ptr<unsigned char[]>, int>  CipherUtils::AES_Encrypt(const char* plainText,
                                                                          int textLength,
                                                                          unsigned char* key,
                                                                          unsigned char* iv) {

    // 创建一个新的加密上下文ctx
    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    // 设置密文的空间
    std::shared_ptr<unsigned char[]> cipherText(new unsigned char[textLength + AES_BLOCK_SIZE]);
    int len;
    int cipherLength = 0;
    // 初始化加密
    EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv);
    // 加密
    EVP_EncryptUpdate(ctx, cipherText.get(), &len, (const unsigned char*)plainText, textLength);
    cipherLength = len;

    // 结束加密操作
    EVP_EncryptFinal_ex(ctx, cipherText.get() + len, &len);
    cipherLength += len;
    // 清理
    EVP_CIPHER_CTX_free(ctx);
    // 返回密文与密文长度
    return { cipherText, cipherLength };
}

/*使用aes加密*/
std::pair<std::shared_ptr<unsigned char[]>, int>  CipherUtils::AES_Encrypt(std::shared_ptr<char[]> plainText,
                                                                          int textLength,
                                                                          std::shared_ptr<unsigned char[]> key,
                                                                          std::shared_ptr<unsigned char[]> iv) {
    return AES_Encrypt(plainText.get(), textLength, key.get(), iv.get());
}

/*使用aes加密*/
std::pair<std::shared_ptr<unsigned char[]>, int>  CipherUtils::AES_Encrypt(const std::string& plainText,
                                                                          std::shared_ptr<unsigned char[]> key,
                                                                          std::shared_ptr<unsigned char[]> iv) {
    return AES_Encrypt(plainText.c_str(), plainText.size(), key.get(), iv.get());
}


/*使用aes解密*/
std::pair<std::shared_ptr<unsigned char[]>, int> CipherUtils::AES_Decrypt(unsigned char* cipherText,
                                                                          int cipherLength,
                                                                          unsigned char* key,
                                                                          unsigned char* iv) {

    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    std::shared_ptr<unsigned char[]> plaintext(new unsigned char[cipherLength]);
    // 初始化成0
    memset(plaintext.get(), 0, cipherLength);
    int len;
    int plainTextLength;
    // 初始化解密操作
    EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv);

    // 解密
    EVP_DecryptUpdate(ctx, plaintext.get(), &len, cipherText, cipherLength);
    plainTextLength = len;

    // 结束解密操作
    EVP_DecryptFinal_ex(ctx, plaintext.get() + len, &len);
    plainTextLength += len;

    // 清理
    EVP_CIPHER_CTX_free(ctx);

    return { plaintext, plainTextLength };
}

/*使用aes解密*/
std::pair<std::shared_ptr<unsigned char[]>, int> CipherUtils::AES_Decrypt(std::shared_ptr<unsigned char[]> cipherText,
                                                                          int cipherLength,
                                                                          std::shared_ptr<unsigned char[]> key,
                                                                          std::shared_ptr<unsigned char[]> iv) {
    return AES_Decrypt(cipherText.get(), cipherLength, key.get(), iv.get());
}

/*生成ecc密钥对，生成pem格式的密钥对*/
std::pair<std::string, std::string> CipherUtils::ECC_generate() {

    std::string publicKey, privateKey;
    EC_KEY* ecKey = EC_KEY_new_by_curve_name(NID_X9_62_prime256v1); // 使用P-256曲线

    if (ecKey != nullptr && EC_KEY_generate_key(ecKey)) {
        BIO* pub = BIO_new(BIO_s_mem());
        BIO* priv = BIO_new(BIO_s_mem());

        if (PEM_write_bio_EC_PUBKEY(pub, ecKey) && PEM_write_bio_ECPrivateKey(priv, ecKey, nullptr, nullptr, 0, nullptr, nullptr)) {
            size_t pubLength = BIO_pending(pub);
            size_t privLength = BIO_pending(priv);

            publicKey.resize(pubLength + 1);
            privateKey.resize(privLength + 1);

            BIO_read(pub, (void *)publicKey.c_str(), pubLength);
            BIO_read(priv, (void*)privateKey.c_str(), privLength);

            publicKey[pubLength] = '\0';
            privateKey[privLength] = '\0';
        }

        BIO_free_all(pub);
        BIO_free_all(priv);
    }

    EC_KEY_free(ecKey);

    return { publicKey, privateKey };
}
/*计算*/
std::pair<std::shared_ptr<unsigned char[]>, int> CipherUtils::ECDH_computeKey(const std::string& privateKey, const std::string& publicKey) {

    EC_KEY* pubKey = ECC_readPublicKeyFromString(publicKey);
    EC_KEY* privKey = ECC_readPrivateKeyFromString(privateKey);

    size_t field_size = EC_GROUP_get_degree(EC_KEY_get0_group(privKey));
    size_t secret_len = (field_size + 7) / 8;

    std::shared_ptr<unsigned char[]> secret(new unsigned char[secret_len]);

    if (!ECDH_compute_key(secret.get(), secret_len, EC_KEY_get0_public_key(pubKey), privKey, NULL)) {

        EC_KEY_free(pubKey);
        EC_KEY_free(privKey);

        throw std::runtime_error("ECDH key computation failed.");
        return {};
    }
    EC_KEY_free(pubKey);
    EC_KEY_free(privKey);
    return { secret , secret_len };
}


// 读取ECC的pem格式的公钥并返回对应的对象
EC_KEY* CipherUtils::ECC_readPublicKeyFromString(const std::string& publicKeyPEM) {
    BIO* bio = BIO_new_mem_buf((void*)publicKeyPEM.c_str(), -1);
    if (bio == NULL) {
        return NULL;
    }

    EC_KEY* publicKey = PEM_read_bio_EC_PUBKEY(bio, NULL, NULL, NULL);
    BIO_free(bio);

    return publicKey;
}

// 读取ECC的pem格式的私钥并返回对应的对象
EC_KEY* CipherUtils::ECC_readPrivateKeyFromString(const std::string& privateKeyPEM) {
    BIO* bio = BIO_new_mem_buf((void*)privateKeyPEM.c_str(), -1);
    if (bio == NULL) {
        return NULL;
    }

    EC_KEY* privateKey = PEM_read_bio_ECPrivateKey(bio, NULL, NULL, NULL);
    BIO_free(bio);

    return privateKey;
}

/*使用ECHD共享密钥计算AES256密钥与iv*/
std::vector<std::shared_ptr<unsigned char[]>> CipherUtils::AES_generateFromECDHSharedKey(std::shared_ptr<unsigned char[]> ECDHKey, int ECDHKeyLength) {
    // 定义AES-256密钥的长度
    size_t AES256_KEY_LENGTH = 32;
    size_t AES256_IV_LENGTH = 16;

    // 创建两个个buffer来保存派生的AES密钥与IV
    std::shared_ptr<unsigned char[]> AESKey(new unsigned char[AES256_KEY_LENGTH]);
    std::shared_ptr<unsigned char[]> AESIv(new unsigned char[AES256_IV_LENGTH]);


    // 创建并初始化HKDF context
    EVP_PKEY_CTX* pctx = EVP_PKEY_CTX_new_id(EVP_PKEY_HKDF, NULL);
    if (!pctx) {
        throw std::runtime_error("EVP_PKEY_CTX_new_id failed.");
    }
    // 设置盐值
    const char* salt = "ghost-him";

    // HKDF提取阶段
    if (EVP_PKEY_derive_init(pctx) <= 0 ||
        EVP_PKEY_CTX_set_hkdf_md(pctx, EVP_sha256()) <= 0 ||
        EVP_PKEY_CTX_set1_hkdf_salt(pctx, (const unsigned char*)salt, strlen(salt)) <= 0 || // 使用salt="ghost-him"
        EVP_PKEY_CTX_set1_hkdf_key(pctx, ECDHKey.get(), ECDHKeyLength) <= 0) {
        EVP_PKEY_CTX_free(pctx);
        throw std::runtime_error("HKDF extract phase failed.");
    }

    const char* info = "yuzusoft";

    // 派生足够长度的数据以供密钥和IV使用
    size_t outlen = AES256_KEY_LENGTH + AES256_IV_LENGTH;
    std::unique_ptr<unsigned char[]> out = std::make_unique<unsigned char[]>(outlen);

    // HKDF扩展阶段
    if (EVP_PKEY_CTX_add1_hkdf_info(pctx, (const unsigned char*)info, strlen(info)) <= 0 ||
        EVP_PKEY_derive(pctx, out.get(), &outlen) <= 0) {
        EVP_PKEY_CTX_free(pctx);
        throw std::runtime_error("HKDF expand phase failed.");
    }

    // 分割输出以获得密钥和IV
    memcpy(AESKey.get(), out.get(), AES256_KEY_LENGTH);
    memcpy(AESIv.get(), out.get() + AES256_KEY_LENGTH, AES256_IV_LENGTH);

    // 清理
    EVP_PKEY_CTX_free(pctx);

    // 返回AES密钥及其长度
    return { AESKey, AESIv };
    
}

std::string CipherUtils::Base64_Encode(const unsigned char *binary, int length)
{
    BIO *bio, *b64;
    BUF_MEM *bufferPtr;

    b64 = BIO_new(BIO_f_base64());
    bio = BIO_new(BIO_s_mem());

    // 将bio和b64链起来
    bio = BIO_push(b64, bio);

    // 设置不要换行
    BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);

    // 写入数据到BIO
    BIO_write(bio, binary, length);
    BIO_flush(bio);

    // 获取内存指针
    BIO_get_mem_ptr(bio, &bufferPtr);

    // 创建一个std::string来保存编码后的数据
    std::string b64text(bufferPtr->data, bufferPtr->length);

    // 由于BIO_get_mem_ptr不会增加引用计数，所以在BIO_free_all之后bufferPtr可能指向无效内存
    // BIO_free_all会释放与bio相关的所有BIOs，包括b64
    BIO_free_all(bio);

    // 返回编码后的数据
    return b64text;
}

std::pair<std::shared_ptr<unsigned char[]>, int> CipherUtils::Base64_Decode(const std::string &base64String)
{
    BIO *bio, *b64;
    std::shared_ptr<unsigned char[]> output(new unsigned char[base64String.size()]);
    int decodedLength = 0;

    bio = BIO_new_mem_buf(base64String.c_str(), base64String.size());
    b64 = BIO_new(BIO_f_base64());
    bio = BIO_push(b64, bio);

    BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL); // 不使用换行符
    decodedLength = BIO_read(bio, output.get(), base64String.size());
    BIO_free_all(bio);
    return {output, decodedLength};
}

size_t CipherUtils::calculateAESEncryptedLength(size_t plainTextLength)
{
    // 加密后的长度是原始长度加上一个块大小（用于IV）
    size_t encryptedLength = plainTextLength + AES_BLOCK_SIZE;

    // 如果原始长度不是块大小的倍数，还需要填充到下一个块边界
    if (plainTextLength % AES_BLOCK_SIZE != 0) {
        encryptedLength += AES_BLOCK_SIZE - (plainTextLength % AES_BLOCK_SIZE);
    }

    return encryptedLength;
}

std::string CipherUtils::SHA256_calculate(const std::string &input)
{
    unsigned char hash[SHA256_DIGEST_LENGTH]; // SHA256输出的哈希值大小
    SHA256_CTX sha256;
    SHA256_Init(&sha256); // 初始化SHA256上下文
    SHA256_Update(&sha256, input.c_str(), input.size()); // 添加数据到SHA256上下文
    SHA256_Final(hash, &sha256); // 完成哈希计算

    // 将哈希值转换为十六进制字符串
    std::stringstream ss;
    for (int i = 0; i < SHA256_DIGEST_LENGTH; i++) {
        ss << std::hex << std::setw(2) << std::setfill('0') << (int)hash[i];
    }

    return ss.str(); // 返回十六进制的哈希值字符串
}
